//-----------------------------------------------------------------------------
//
// @file	Octree.cpp
// @brief	Octree file
// @author	Vinod Melapudi
// @date	13Aug2011
//
//-----------------------------------------------------------------------------

#include "Octree.h"
#include "MemPool.h"
#include "Utils.h"
#include "assert.h"

using namespace std;



template<class T>
Octree<T>::Octree( uint size )
:m_size(size)
,m_root(NULL)
,m_branch_pool(NULL)
,m_leaf_pool(NULL)
{
	// Check if size is power of 2 and size is greater than 1
	assert( IsPowerOf2(size) );
	assert(size>1);

	// Initiate iterator cache
	m_end.Set( this, NULL, uLocation(size,size,size) );
	m_begin = m_end;
}


template<class T>
Octree<T>::~Octree()
{
	// Delete pool
	delete m_branch_pool;
	delete m_leaf_pool;
	m_root = NULL;
	m_branch_pool = NULL;
	m_leaf_pool = NULL;
}






// Helper fn to get the child index of a location at a given height
// Heigh of the tree is given as the size of number of childrens
inline short GetChildIdx( const uLocation& l, uint size )
{
	// Convert x to  0 if x < size   &   1 if x > size,    same for y,z
	short ix = !!(l.x & size);
	short iy = !!(l.y & size);
	short iz = !!(l.z & size );

	// Get child index
	return (ix << 2 | iy << 1 | iz);
}




// Set a data pointer to a leaf in octree
template<class T>
void	Octree<T>::Set( const uLocation& l, T data )
{
	assert( l.x < m_size && l.y < m_size && l.z < m_size );


	// Create mempool for branch nodes & leaf nodes
	if ( !m_root )
	{
		// Height of the tree = h
		// Max leaves in one axis =  2^(h-1),   so  max leaves in tree = (2^(h-1))^3 = 2^((h-1)*3) = 8^(h-1)
		// so  h  = log2( num_leaves_per_axis ) + 1  =  1 bit location in a num
		uint height_of_tree = 0;
		uint size = m_size;
		while ( size )
		{
			size = size >> 1;
			height_of_tree++;
		}

		// When pool runs out of allocations, will be incremented by height of tree-1
		m_branch_pool = new MemPool( sizeof(Branch), 0, height_of_tree-1 );

		m_leaf_pool = new MemPool( sizeof(Leaf), 0, 8 );
	}


	Node** n = &m_root;
	uint size = m_size;

	// Traverse to the leaf
	while ( size-1 > 0 )
	{
		// If there is no node at this level, create a branch
		Branch* branch;
		if ( *n == NULL )
		{
			branch = (Branch*) m_branch_pool->Allocate( sizeof(Branch) );
			new(branch) Branch();
			*n = branch;
		}
		else
		{
			branch = reinterpret_cast<Branch*>(*n);
		}


		// Divide size
		size = size >> 1;

		// Get child index at this height and get the pointer to child node
		short child_idx = GetChildIdx( l, size );
		n = &branch->m_children[0][0][0]+child_idx;
	}


	// Set leaf
	Leaf* leaf;
	if ( *n )
	{
		leaf = reinterpret_cast<Leaf*>(*n);
	}
	else
	{
		// Create new leaf
		leaf = (Leaf*) m_leaf_pool->Allocate( sizeof(Leaf) );
		new(leaf) Leaf();
		*n = leaf;

		// Reset begin
		m_begin.Location().Set(0,0,0);
		FindLeafPathAt( m_root, m_begin, true );
	}

	// Set data
	leaf->m_data = data;
}


template<class T>
T*	Octree<T>::Get( const uLocation& l ) const
{
	if ( !(	m_root && 
			l.x >= 0 && l.y >= 0 && l.z >= 0 &&
			l.x < m_size && l.y < m_size && l.z < m_size) )
		return NULL;


	Node* n = m_root;
	uint size = m_size;
	
	// Traverse to the leaf (get path & leaf node)
	while ( size-1 > 0 )
	{
		if ( n == NULL )
			return NULL;

		// Get branch
		Branch* branch = reinterpret_cast<Branch*>(n);

		// Divide size
		size = size >> 1;

		// Get child index at this height and get then the child node
		short child_idx = GetChildIdx( l, size );
		n = branch->Child(child_idx);
	}


	// Get leaf & return the iterator
	if ( n )
	{
		Leaf* leaf = reinterpret_cast<Leaf*>(n);
		return &leaf->m_data;
	}

	return NULL;
}

template<class T>
T*	Octree<T>::operator()( const uLocation& l ) const
{
	return Get(l);
}


// Get iterator to a given location (if location doesnt exist, return end())
template<class T>
typename Octree<T>::iterator	Octree<T>::Iter( const uLocation& l ) const
{
	T* dataptr = Get(l);
	if ( dataptr )
	{
		iterator iter;
		iter.Set( this, dataptr, l );
		return iter;
	}

	return m_end;
}




// Get list of branches the current iterator points to
template<class T>
void	Octree<T>::GetPathBranches( iterator& iter, PathBranches& path_branches ) const
{
	assert( m_root );

	const uLocation l = iter.Location();

	Node* n = m_root;
	uint size = m_size;
	
	// Traverse to the leaf & get the path branches
	while ( size-1 > 0 )
	{
		if ( n == NULL )
			return;

		// Get branch
		Branch* branch = reinterpret_cast<Branch*>(n);
		path_branches.push_back( branch );

		// Divide size
		size = size >> 1;

		// Get child index at this height and get then the child node
		short child_idx = GetChildIdx( l, size );
		n = branch->Child(child_idx);
	}
}

// Given a iterator, update it with the path to next leaf in the given direction
template<class T>
void	Octree<T>::UpdateIterator( iterator& iter, bool positive_dir ) const
{
	// If iterator is pointing to first, block iterating in negative dir
	if ( iter == m_begin && !positive_dir  )
	{
		return;
	}

	// If iterator is pointing to last, then block iterating in positive dir
	if ( iter == m_end  )
	{
		if ( positive_dir )
			return;

		iter.Location().Set(0,0,0);
		FindLeafPathAt( m_root, iter, positive_dir );
		return;
	}


	// Get list of branch nodes the current iterator points to
	PathBranches path_branches;
	GetPathBranches( iter, path_branches );


	int addby = positive_dir ? 1 : -1;

	// Iterate from bottom of the path, looking for next leaf
	for( int i=(int)path_branches.size()-1; i >= 0; i-- )
	{
		// Get current child index
		const uLocation& l = iter.Location();
		short child_idx = ((l.x & 1) << 2) | ((l.y & 1) << 1) | (l.z & 1);

		// Get next child to find leaf at
		short next_child_idx = child_idx + addby;

		// Pop current path
		iter.PopLocation();

		// Find new path with remaining children
		while( next_child_idx >=0 && next_child_idx <= 7 )
		{
			// Get child
			Node* n = path_branches[i]->Child( next_child_idx );

			// Push current child
			iter.PushLocation( next_child_idx );

			// Find leaf in current child
			if ( FindLeafPathAt( n, iter, positive_dir ) )
				return;

			// Move to next child
			iter.PopLocation();
			next_child_idx += addby;
		}
	}

	// If we reached here, no leaf was found
	iter = m_end;
}




// Recursively iterate down a given branch node, to the 1st or last leaf node, 
// iterator's location should match node thats passed at the same level, 
// so if node points to root, then location should be (0,0,0)
// So number of bits in location is same as height of the node in the tree
template<class T>
bool	Octree<T>::FindLeafPathAt( typename Octree<T>::Node* node, iterator& iter, bool positive_dir ) const
{
	if ( !node )
		return false;
	
	// If node is a leaf, then return true
	if ( node->leaf )
	{
		// Fill iterator info
		Leaf* leaf = reinterpret_cast<Leaf*>(node);
		iter.Set( this, &leaf->m_data );
		return true;
	}

	// Check all children of the branch
	Branch* branch = reinterpret_cast<Branch*>(node);
	for( short i=0; i < 8; i++ )
	{
		// Reverse index if searching for last leaf
		short index = positive_dir ? i : 7-i;

		// Get child
		Node* child = branch->Child(index);

		// Construct location based on index
		iter.PushLocation(index);

		if ( FindLeafPathAt( child, iter, positive_dir ) )
			return true;

		// Re-construct location, if no leaf was found at this level
		iter.PopLocation();
	}

	return false;
}





// Increment iterator
template<class T>
typename Octree<T>::iterator&	Octree<T>::iterator::operator++()
{
	// Get next leaf along the path
	m_octree->UpdateIterator( *this, true );
	return *this;
}

// Decrement iterator
template<class T>
typename Octree<T>::iterator&	Octree<T>::iterator::operator--()
{
	// Get next leaf along the path
	m_octree->UpdateIterator( *this, false );
	return *this;
}

// Compare locations of 2 leaves
template<class T>
bool	Octree<T>::iterator::operator<( const typename Octree<T>::iterator&	rhs ) const
{
	if ( m_location.x != rhs.m_location.x )
		return m_location.x < rhs.m_location.x;
	
	if( m_location.y != rhs.m_location.y )
		return m_location.y < rhs.m_location.y;
	
	return m_location.z < rhs.m_location.z;
}

// Compare locations of 2 leaves
template<class T>
bool	Octree<T>::iterator::operator==( const typename Octree<T>::iterator& rhs ) const
{
	return m_location == rhs.m_location;
}

// Push an index to child node to the location
template<class T>
void	Octree<T>::iterator::PushLocation( short idx )
{
	m_location.x = (m_location.x << 1) | !!(idx & 4);
	m_location.y = (m_location.y << 1) | !!(idx & 2);
	m_location.z = (m_location.z << 1) | !!(idx & 1);
}

// Pop the location back to parent node
template<class T>
void	Octree<T>::iterator::PopLocation()
{
	m_location.x = m_location.x >> 1;
	m_location.y = m_location.y >> 1;
	m_location.z = m_location.z >> 1;
}




// Testing octree
static bool OctreeTest()
{/*
	{
		// Check 1 level octree
		Octree<float>* o = new Octree<float>(2);
		
		o->Set( uLocation(0,0,0), 1);
		o->Set( uLocation(0,0,1), 2);
		o->Set( uLocation(0,1,0), 3);
		o->Set( uLocation(1,0,0), 4);
		o->Set( uLocation(0,0,0), 1);
		o->Set( uLocation(0,0,1), 2);

		float* const t = o->Get( uLocation(0,0,1) );
		assert( (*t) == 2.f );
		(*t) = 8.0;

		// Iterators
		Octree<float>::iterator biter = o->Begin();
		Octree<float>::iterator eiter = o->End();
		float* s = (*biter);


		bool bb = biter < eiter;
		biter++;
		biter++;
		biter++;
		biter++;

		assert( biter == eiter );
	
		biter--;
		biter--;
		biter--;
		biter--;

		assert( biter < eiter );
		assert( biter == o->Begin() );

		delete o;
	}
	*/

	{
		// Check 2 level octree
		Octree<float>* o = new Octree<float>(4);

		o->Set( uLocation(0,0,0), 1);
		o->Set( uLocation(0,2,1), 2);
		o->Set( uLocation(3,1,0), 3);
		o->Set( uLocation(1,0,3), 4);
		o->Set( uLocation(2,1,2), 5);
		o->Set( uLocation(3,0,1), 6);

		float* const t1 = o->Get( uLocation(2,0,1) );
		assert( t1 == NULL );

		float* const t2 = o->Get( uLocation(2,1,2) );
		assert( (*t2) == 5.f );
		(*t2) = 8.0;

		Octree<float>::iterator iter = o->Iter( uLocation(1,0,3) );
		assert( (*(*iter)) == 4.f );

		// Iterators
		Octree<float>::iterator biter = o->Begin();
		Octree<float>::iterator eiter = o->End();
		float* s = (*biter);


		bool bb = biter < eiter;
		
		biter++;
		biter++;
		biter++;
		biter++;
		biter++;
		biter++;
		biter++;
		biter++;

		assert( biter == eiter );

		biter--;
		biter--;
		biter--;
		biter--;
		biter--;
		biter--;
		biter--;
		biter--;

		assert( biter == o->Begin() );

		delete o;
	}

	return true;
}